#include <kutils.h>
#include "paging.h"
#include "memory.h"
#include "display.h"
#include "keyboard.h"

s32int page_create_table(address_space_t *space, u32int table_index);
s32int init_address_space(address_space_t* space);

const u32int page_size = 0x1000;
const u32int table_space_start = 0x0FC00000;
const u32int table_space_end = 0x0FFFFFFF;

address_space_t* kspace;
address_space_t* active_space;

extern u32int end;
extern u32int free_mem_start;
extern u32int mem_features;

s32int page_create_table(address_space_t *space, u32int table_index)
{
	if(!(mem_features & MEM_FEATURES_VIRT))
		return 0;

#ifdef FLUIDIUM_DEBUG
	write_string("Memory: Create Table: Space: ");
	write_hex((u32int)space);
	write_string(" Index: ");
	write_number(table_index);
	write_string("\n");
#endif

	//u32int page_index = table_index * 1024;
	u32int address = phys_alloc(1024);
	if(!address)
		return 0;
	space->page_directory[table_index] = address; 
	space->page_directory[table_index] |= 3;

	space->page_tables[63][table_index] = address | 3;
	space->page_tables[table_index] = (u32int*)(0x0FC00000 + (table_index * page_size));

	s32int i = 0;
	for(i = 0; i < 1024; i++)
	{
		space->page_tables[table_index][i] = 0;
	}
	return 1;
}

s32int init_address_space(address_space_t* space)
{
#ifdef FLUIDIUM_DEBUG
	write_string("Memory: Init Address Space: ");
	write_hex((u32int)space);
	write_string("\n");
#endif

	//create all the page directories
	if(mem_features & MEM_FEATURES_VIRT)
	{
		space->page_directory = (u32int*)id_page_alloc(1024);
	}
	else
	{
		space->page_directory = (u32int*)phys_alloc(1024);
	}
	if(!space->page_directory)
		return 0;
	s32int i = 0;
	for(i = 0; i < 1024; i++)
	{
		// attribute set to: supervisor level, read/write, not present(010 in binary)
		space->page_directory[i] = 0;
	}

	//declare the page table for this address space where all of the page tables will be mapped.
	//The magic number 63 is the index of the page table that maps the addresses from 
	//0x0FC00000 to 0x0FFFFFFF.
	u32int address;
	if(mem_features & MEM_FEATURES_VIRT)
	{
		address = id_page_alloc(1024);
	}
	else
	{
		address = phys_alloc(1024);
	}
	if(!address)
		return 0;
	u32int page_index = 0x0FC00000 / page_size;
	u32int table_index = 63;
	space->page_tables[table_index] = (u32int*)address;
	for(i = 0; i < 1024; i++)
	{
		space->page_tables[table_index][i] = 0;
	}
	space->page_directory[table_index] = address; 
	space->page_directory[table_index] |= 3;

	//map this page table s32into itself
	//any access to the address 0x0FC0FC00 will now be translated to the physical address of the page table.
	space->page_tables[table_index][page_index - (table_index * 1024)] = address | 3; 
	
	//declare the first page table which will map the first 4 megabytes before paging is enabled.
	page_index = 0x00 / page_size;
	if(mem_features & MEM_FEATURES_VIRT)
	{
		address = id_page_alloc(1024);
	}
	else
	{
		address = phys_alloc(1024);
	}
	if(!address)
		return 0;
	space->page_tables[0] = (u32int*)address;
	for(i = 0; i < 1024; i++)
	{
		space->page_tables[0][i] = 0;
	}
	space->page_directory[0] = address; 
	space->page_directory[0] |= 3;

	space->page_tables[63][0] = address | 3; 
	
	u32int page_index_start = 0 / page_size;
	u32int page_index_end = (free_mem_start / page_size) + 1;
	page_index = page_index_start;
	address = 0;
	for(page_index = page_index_start; page_index < page_index_end; page_index++)
	{
		table_index = page_index / 1024;

		space->page_tables[table_index][page_index - (table_index * 1024)] = address | 3; 
		address += page_size;
	}

	//the page table should now be accessed through it's new virtual address instead of it's physical one.
	space->page_tables[table_index] = (u32int*)(0x0FC00000 + (63 * page_size));

	//the page table should now be accessed through it's new virtual address instead of it's physical one.
	space->page_tables[0] = (u32int*)(0x0FC00000);
	return 1;
}

s32int switch_address_space(address_space_t* space)
{
	active_space = space;
	u32int cr0;
	__asm__ __volatile__("mov %0, %%cr3":: "b"(active_space->page_directory));
   __asm__ __volatile__("mov %%cr0, %0": "=b"(cr0));
   cr0 |= 0x80000000;
   __asm__ __volatile__("mov %0, %%cr0":: "b"(cr0));
	return 1;
}

void disable_paging(void)
{
	u32int cr0;
	__asm__ __volatile__("mov %%cr0, %0": "=b"(cr0));
   cr0 &= ~0x80000000;
   __asm__ __volatile__("mov %0, %%cr0":: "b"(cr0));
}

s32int map_address_range(address_space_t *space, u32int virt_addr_start, u32int virt_addr_end, u32int phys_addr)
{
	if(!(mem_features & MEM_FEATURES_VIRT))
		return 0;

	//convert the virtual address s32into a page index
	u32int page_index_start = virt_addr_start / page_size;
	u32int page_index_end = (virt_addr_end / page_size) + 1;
	u32int page_index = page_index_start;
	u32int address = phys_addr;
	u32int table_index = 0;

	for(page_index = page_index_start; page_index < page_index_end; page_index++)
	{
		table_index = page_index / 1024;
		if(!space->page_tables[table_index])
		{
			if(!page_create_table(space, table_index))
				return 0;
		}

		space->page_tables[table_index][page_index - (table_index * 1024)] = address | 3; 
		address += page_size;
	}
	return 1;
}

s32int page_alloc(address_space_t *space, u32int virt_addr_start, u32int virt_addr_end)
{
	if(!(mem_features & MEM_FEATURES_VIRT))
		return 0;

	u32int tmp = phys_alloc(virt_addr_end - virt_addr_start);
	if(!tmp)
		return 0;
	if(!map_address_range(space, virt_addr_start, virt_addr_end, tmp))
		return 0;
	return 1;
}

u32int id_page_alloc(u32int size)
{
	if(!(mem_features & MEM_FEATURES_VIRT))
		return 0;

	u32int addr = phys_alloc(size);
	map_address_range(active_space,addr,addr + size, addr);
	return addr;
}

s32int page_free(address_space_t* space, u32int virt_addr_start, u32int virt_addr_end)
{
	if(!(mem_features & MEM_FEATURES_VIRT))
		return 0;

	//convert the virtual address s32into a page index
	u32int page_index_start = virt_addr_start / page_size;
	u32int page_index_end = (virt_addr_end / page_size);
	u32int page_index = page_index_start;
	u32int addr_phys;
	u32int table_index;

	for(page_index = page_index_start; page_index < page_index_end; page_index++)
	{
		table_index = page_index / 1024;

		if(space->page_tables[table_index])//only free from tables that actually exist
		{
			addr_phys = space->page_tables[table_index][page_index - (table_index * 1024)] & 0xFFFFF000;

			//unmap the virtual address
			space->page_tables[table_index][page_index - (table_index * 1024)] = 0;
					
			//free the physical address
			if(!phys_free(addr_phys))
				return 0;
		}
	}
	return 1;
}

address_space_t* create_address_space(void)
{
	address_space_t* dest = 0;
	if(mem_features & MEM_FEATURES_VIRT)
	{
		dest = (address_space_t*)id_page_alloc(sizeof(address_space_t));
	}
	else
	{
		dest = (address_space_t*)phys_alloc(sizeof(address_space_t));
	}
	if(!dest)
		return 0;
	memset(dest, 0, sizeof(address_space_t));//important

	if(!init_address_space(dest))
		return 0;

	// NOTE:
	// In the kernel heap area the pos32inter to the table should be cloned insted of each page individually

	if(kspace)
	{
		u32int table_index = 0;
		u32int page_index = 0;
		for(table_index = 0; table_index < 256; table_index++)
		{
			if(table_index == 63)//don't clone the table s32into which the table itself is mapped.
				continue;
			if(kspace->page_tables[table_index])//if the table exists in the kernel space...
			{	// create it in the new space
				//if(table_index < 256) // if this table is a part of kernel space
				//{	// clone whole tables
				//	u32int address = (u32int)kspace->page_tables[table_index];
				//	write_hex((u32int)address);
				//	halt();
				//}
				//else
				//{	// clone each page individually
					u32int address = id_page_alloc(1024);
					dest->page_directory[table_index] = address; 
					dest->page_directory[table_index] |= 3;
					dest->page_tables[63][table_index] = address | 3;
					dest->page_tables[table_index] = (u32int*)address;//temporarially pos32int the table to the address so it can be filled
				
					for(page_index = 0; page_index < 1024; page_index++)
					{
						if(kspace->page_tables[table_index][page_index])
						{
							dest->page_tables[table_index][page_index] = 
								kspace->page_tables[table_index][page_index];
							//write_hex(dest->page_tables[table_index][page_index]);
						}
					}
					dest->page_tables[table_index] = (u32int*)(0x0FC00000 + (table_index * page_size));
				//}
			}
		}
	}
	return dest;
}

s32int destroy_address_space(address_space_t* space)
{
	if(!space)
		return 0;

	u32int i = 0;
	for(i = 0; i < 1024; i++)
	{
		if(space->page_directory[i])
			if(!phys_free(space->page_directory[i]))
				return 0;
	}
	if(!phys_free((u32int)space))
		return 0;
	return 1;
}

void initialize_paging(void)
{
	kspace = create_address_space();
	switch_address_space(kspace);
	mem_features |= MEM_FEATURES_VIRT;
}

